package uy.edu.fing.pgsegesb.web;

import javax.xml.crypto.*;
import javax.xml.crypto.dsig.*;
import javax.xml.crypto.dom.*;
import javax.xml.crypto.dsig.dom.DOMSignContext;
import javax.xml.crypto.dsig.dom.DOMValidateContext;
import javax.xml.crypto.dsig.keyinfo.*;
import javax.xml.crypto.dsig.spec.C14NMethodParameterSpec;
import javax.xml.crypto.dsig.spec.TransformParameterSpec;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.OutputStream;
import java.security.*;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;

import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.transform.*;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;
import org.w3c.dom.Document;
import org.w3c.dom.NodeList;

public class XMLDSignatureHelper {

	
	public static String sign(String xmlMessage, String elementToSign, String digestMethod, String signatureMethod) throws Exception {

		String keyAlg = "";
		if (signatureMethod.compareTo(SignatureMethod.DSA_SHA1)==0)
			keyAlg = "DSA";
		else
			if(signatureMethod.compareTo(SignatureMethod.RSA_SHA1)==0)
				keyAlg = "RSA";
			else
				return "KeyAlgorithm not supported: only SignatureMethod.DSA_SHA1 or SignatureMethod.RSA_SHA1 available";
		String xmlSignedMessage;

	    // Create a DOM XMLSignatureFactory that will be used to generate the 
		// enveloped signature
	        String providerName = System.getProperty
	            ("jsr105Provider", "org.jcp.xml.dsig.internal.dom.XMLDSigRI");
		XMLSignatureFactory fac = XMLSignatureFactory.getInstance("DOM",
		     (Provider) Class.forName(providerName).newInstance());

	        // Create a Reference to the enveloped document (in this case we are
		// signing the whole document, so a URI of "" signifies that) and
		// also specify the SHA1 digest algorithm and the ENVELOPED Transform.
		Reference ref = fac.newReference
		    ("", fac.newDigestMethod(digestMethod, null),
	             Collections.singletonList(fac.newTransform(Transform.ENVELOPED, (TransformParameterSpec) null)), 
		     null, null);

		// Create the SignedInfo
		SignedInfo si = fac.newSignedInfo
		    (fac.newCanonicalizationMethod
		     (CanonicalizationMethod.INCLUSIVE_WITH_COMMENTS, 
		      (C14NMethodParameterSpec) null), 
		     fac.newSignatureMethod(signatureMethod, null),
		     Collections.singletonList(ref));

	        // Create a DSA KeyPair
	        KeyPairGenerator kpg = KeyPairGenerator.getInstance(keyAlg);
		kpg.initialize(512);
	        KeyPair kp = kpg.generateKeyPair();

	        // Create a KeyValue containing the DSA PublicKey that was generated
		KeyInfoFactory kif = fac.getKeyInfoFactory();
	        KeyValue kv = kif.newKeyValue(kp.getPublic());

		// Create a KeyInfo and add the KeyValue to it
	        KeyInfo ki = kif.newKeyInfo(Collections.singletonList(kv));

		// Instantiate the document to be signed
		DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
		dbf.setNamespaceAware(true);
		Document doc = 
		    dbf.newDocumentBuilder().parse(new ByteArrayInputStream(xmlMessage.getBytes("UTF-8")));

	        // Create a DOMSignContext and specify the DSA PrivateKey and
	        // location of the resulting XMLSignature's parent element
		DOMSignContext dsc = new DOMSignContext
		    (kp.getPrivate(), doc.getDocumentElement());

		// Create the XMLSignature (but don't sign it yet)
		XMLSignature signature = fac.newXMLSignature(si, ki);

	        // Marshal, generate (and sign) the enveloped signature
	        signature.sign(dsc);

		// output the resulting document
		OutputStream os = new ByteArrayOutputStream();

    	TransformerFactory tf = TransformerFactory.newInstance();
    	Transformer trans = tf.newTransformer();
    	trans.transform(new DOMSource(doc), new StreamResult(os));
    	
    	xmlSignedMessage = os.toString();
    	return xmlSignedMessage;
    	
    }
     


    public static boolean validate(String xmlMessage) throws Exception {

		// Instantiate the document to be validated
		DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
		dbf.setNamespaceAware(true);
		Document doc = dbf.newDocumentBuilder().parse(new ByteArrayInputStream(xmlMessage.getBytes("UTF-8")));
	
		// Find Signature element
		NodeList nl = 
		    doc.getElementsByTagNameNS(XMLSignature.XMLNS, "Signature");
		if (nl.getLength() == 0) {
		    throw new Exception("Cannot find Signature element");
		}
	
		// Create a DOM XMLSignatureFactory that will be used to unmarshal the 
		// document containing the XMLSignature 
		String providerName = System.getProperty
	            ("jsr105Provider", "org.jcp.xml.dsig.internal.dom.XMLDSigRI");
		XMLSignatureFactory fac = XMLSignatureFactory.getInstance("DOM",
	            (Provider) Class.forName(providerName).newInstance());
	
		// Create a DOMValidateContext and specify a KeyValue KeySelector
	        // and document context
		DOMValidateContext valContext = new DOMValidateContext
		    (new KeyValueKeySelector(), nl.item(0));
		
		// unmarshal the XMLSignature
		XMLSignature signature = fac.unmarshalXMLSignature(valContext);
	
		// Validate the XMLSignature (generated above)
		boolean coreValidity = signature.validate(valContext); 
	
		// Check core validation status
		if (coreValidity == false) {
	    	    System.err.println("Signature failed core validation"); 
		    boolean sv = signature.getSignatureValue().validate(valContext);
		    System.out.println("signature validation status: " + sv);
		    // check the validation status of each Reference
		    Iterator i = signature.getSignedInfo().getReferences().iterator();
		    for (int j=0; i.hasNext(); j++) {
			boolean refValid = 
			    ((Reference) i.next()).validate(valContext);
			System.out.println("ref["+j+"] validity status: " + refValid);
		    }
		} else {
	    	    System.out.println("Signature passed core validation");
		}
		return coreValidity;
    }
    
    private static class KeyValueKeySelector extends KeySelector {
	public KeySelectorResult select(KeyInfo keyInfo,
                                        KeySelector.Purpose purpose,
                                        AlgorithmMethod method,
                                        XMLCryptoContext context)
            throws KeySelectorException {
            if (keyInfo == null) {
		throw new KeySelectorException("Null KeyInfo object!");
            }
            SignatureMethod sm = (SignatureMethod) method;
            List list = keyInfo.getContent();

            for (int i = 0; i < list.size(); i++) {
		XMLStructure xmlStructure = (XMLStructure) list.get(i);
            	if (xmlStructure instanceof KeyValue) {
                    PublicKey pk = null;
                    try {
                        pk = ((KeyValue)xmlStructure).getPublicKey();
                    } catch (KeyException ke) {
                        throw new KeySelectorException(ke);
                    }
                    // make sure algorithm is compatible with method
                    if (algEquals(sm.getAlgorithm(), pk.getAlgorithm())) {
                        return new SimpleKeySelectorResult(pk);
                    }
		}
            }
            throw new KeySelectorException("No KeyValue element found!");
	}

        //@@@FIXME: this should also work for key types other than DSA/RSA
	static boolean algEquals(String algURI, String algName) {
            if (algName.equalsIgnoreCase("DSA") &&
		algURI.equalsIgnoreCase(SignatureMethod.DSA_SHA1)) {
		return true;
            } else if (algName.equalsIgnoreCase("RSA") &&
                       algURI.equalsIgnoreCase(SignatureMethod.RSA_SHA1)) {
		return true;
            } else {
		return false;
            }
	}
    }

    private static class SimpleKeySelectorResult implements KeySelectorResult {
	private PublicKey pk;
	SimpleKeySelectorResult(PublicKey pk) {
	    this.pk = pk;
	}

	public Key getKey() { return pk; }
    }
}
